#!/usr/bin/env python
# -*- coding: utf-8 -*-
#  @Time    : 2021-02-09 15:09
#  @Author  : lifan
#  @File    : test.py
#  @Software: PyCharm
# @Brief   :

import sys

sys.path.append("..")
import numpy as np
import os
import random
import time
from torchvision import transforms
from dataset import MyDataset
from torch.utils.data import DataLoader
from tqdm import tqdm
from model.unet_model import UNet
from loss.loss import *
import multiprocessing
from metric.metric import *
import cv2
from PIL import Image
from ternausnet.models import UNet16
CFG = {
    'fold_num': 1,
    'seed': 1,
    'epochs': 100,
    'images_size': 224,
    'valid_bs': 1,
    'num_workers': multiprocessing.cpu_count(),
    'device': 'cuda:0',
}

# config
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
device = torch.device(CFG['device'])


# Fixed seed
def set_seed(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = True


def test(model, val_dataloader, save_path):
    model.eval()
    best_threshold = 0.5
    with torch.no_grad():
        for step, val_data in enumerate(tqdm(val_dataloader)):
            images, labels, label_path = val_data
            images, labels = images.to(device), labels.to(device)

            save_name = os.path.split(label_path[0])[1].split(".")[0]

            # preds = torch.sigmoid(model(images))
            # preds = preds.detach().cpu().numpy()[:, 0, :, :] # (batch_size, 1, size, size) -> (batch_size, size, size)
            # for probability in preds:
            #     # if probability.shape != (320, 140):
            #     #     probability = cv2.resize(probability, dsize=(320, 140), interpolation=cv2.INTER_LINEAR)
            #     mask = probability > best_threshold
            #     mask = (mask*255).astype(np.uint8)
            #
            #     cv2.imwrite("{}/{}.bmp".format(save_path, str(save_name)), mask)

            # preds = preds.squeeze(0)
            # # preds = preds.detach().cpu().numpy()[:, 0, :, :]  # (batch_size, 1, size, size) -> (batch_size, size, size)
            #
            # tf = transforms.Compose(
            #     [
            #         transforms.ToPILImage(),
            #         transforms.Resize((140,320)),
            #         transforms.ToTensor()
            #     ]
            # )
            # preds = tf(preds.cpu())
            # full_mask = preds.squeeze().cpu().numpy()
            # mask = full_mask > best_threshold
            # mask = (mask*255).astype(np.uint8)
            # # mask = Image.fromarray((mask * 255).astype(np.uint8))

            # cv2.imwrite("{}/{}.bmp".format(save_path, str(save_name)), mask)


            labels = labels.squeeze().cpu().numpy()
            mask = (labels*255).astype(np.uint8)
            cv2.imwrite("{}/{}.bmp".format(save_path, str(save_name)), mask)


if __name__ == '__main__':
    print(time.asctime(time.localtime(time.time())), "Start Training!!!")

    set_seed(CFG['seed'])
    model_weight_path = "/home/fanli/ci2p_dl_code/YingyingHan_UNetSeg_Experiment/natural_images/Trainer/weights/unet_weights_{}.pth".format(
        CFG['fold_num'])
    # model = UNet().to(device)
    model = UNet16(num_classes=1, num_filters=32, pretrained=False).to(device)
    model.load_state_dict(torch.load(model_weight_path))

    val_datasets = MyDataset(fold=CFG['fold_num'], mode="val", images_size=CFG['images_size'])
    val_dataloader = DataLoader(val_datasets, batch_size=CFG['valid_bs'], shuffle=False,
                                num_workers=CFG['num_workers'], pin_memory=torch.cuda.is_available())

    if os.path.exists("./preds_image/fold_{}".format(CFG['fold_num'])) is False:
        os.mkdir("./preds_image/fold_{}".format(CFG['fold_num']))
    save_path = "./preds_image/fold_{}".format(CFG['fold_num'])

    test(model, val_dataloader, save_path)
